import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas


class TSPAnimation:
    def __init__(self, canvas, num_cities=10):
        # 初始化类时生成随机城市坐标
        self.num_cities = num_cities
        self.cities = np.random.rand(num_cities, 2)
        # 计算距离矩阵
        self.distance_matrix = self.calc_distance_matrix(self.cities)
        # 使用最近邻算法求解 TSP
        self.path = self.nearest_neighbor_algorithm(self.distance_matrix)
        # 验证路径是否有效
        self.is_valid_path = self.verify_path(self.path, self.num_cities)
        # 计算路径总距离
        self.total_distance = self.calculate_total_distance(self.path, self.distance_matrix)
        # 生成验证信息和路径总距离的文本
        self.validation_text = "Path validation: PASSED" if self.is_valid_path else "Path validation: FAILED"
        self.distance_text = f"Total distance: {self.total_distance:.2f}"


        self.canvas = canvas

        # 清除上一次绘制的内容
        self.canvas.figure.clear()
        self.canvas.draw()


        self.fig, self.ax = plt.subplots()
        self.canvas.figure = self.fig
        self.canvas.draw()


    def calc_distance_matrix(self, cities):
        # 计算每对城市之间的欧几里得距离，并存储在 distance_matrix 中
        num_cities = len(cities)
        distance_matrix = np.zeros((num_cities, num_cities))
        for i in range(num_cities):
            for j in range(num_cities):
                distance_matrix[i, j] = np.linalg.norm(cities[i] - cities[j])
        return distance_matrix

    def nearest_neighbor_algorithm(self, distance_matrix):
        # 使用最近邻算法找到一个近似的最短路径
        num_cities = len(distance_matrix)
        visited = [False] * num_cities
        path = [0]
        visited[0] = True
        for _ in range(num_cities - 1):
            last_city = path[-1]
            next_city = np.argmin(
                [distance_matrix[last_city, j] if not visited[j] else np.inf for j in range(num_cities)])
            path.append(next_city)
            visited[next_city] = True
        path.append(0)
        return path

    def verify_path(self, path, num_cities):
        # 检查路径是否有效，即包含所有城市，并且起点和终点相同
        if len(path) != num_cities + 1:
            return False
        if len(set(path[:-1])) != num_cities:
            return False
        if path[0] != path[-1]:
            return False
        return True

    def calculate_total_distance(self, path, distance_matrix):
        # 计算路径的总距离
        total_distance = 0
        for i in range(len(path) - 1):
            total_distance += distance_matrix[path[i], path[i + 1]]
        return total_distance

    def begin_animation(self):
        # 初始化绘图
        self.ax.clear()
        scat = self.ax.scatter(self.cities[:, 0], self.cities[:, 1], c='blue')
        self.line, = self.ax.plot([], [], 'r-', lw=2)
        self.step_text = self.ax.text(0.02, 0.95, '', transform=self.ax.transAxes)


        # 在图形外部显示验证信息和路径总距离
        self.fig.text(0.02, 0.95, self.validation_text, transform=self.fig.transFigure)
        self.fig.text(0.02, 0.92, self.distance_text, transform=self.fig.transFigure)

        # 标出起始城市
        start_city = self.cities[0]
        self.ax.scatter(start_city[0], start_city[1], c='green', s=200)

        # 显示每个城市的坐标
        for i, (x, y) in enumerate(self.cities):
            self.ax.text(x, y, f'({x:.2f}, {y:.2f})', fontsize=9, ha='right')

        # 创建动画
        self.ani = animation.FuncAnimation(self.fig, self.update, frames=len(self.path) - 1, init_func=self.init,
                                           blit=True, repeat=False, interval=500)
        self.canvas.draw()

    def init(self):
        # 清空路径线和步骤文本
        self.line.set_data([], [])
        self.step_text.set_text('')
        return self.line, self.step_text

    def update(self, frame):
        # 更新路径线和步骤文本以反映当前的路径进度
        step = frame + 1
        current_path = self.path[:step]
        x = [self.cities[i][0] for i in current_path]
        y = [self.cities[i][1] for i in current_path]
        self.line.set_data(x, y)
        self.step_text.set_text(f'Step: {step}')
        return self.line, self.step_text


